/*
* Description: 基于etw的应用使用统计
*/
#include "app_using_statistics.h"
#include "model/app_using_data.h"
#include "log.h"
#include "text_util/string_proc.h"
#include <thread>
#include <mutex>
#include "process_load_statistic/payload_sensor.h"
#include "process_detect.h"
#include "performance/perf_event_dispatcher.h"
#include "app_controller/app_controller.h"
#include "performance/install_lifecycle_handler.h"
namespace DFX {

const int APP_CREATE_EVENT_FILE_PATH_OFFSET = 60;

const int APP_QUIT_EVENT_FILE_PATH_OFFSET = 84;

const int CLIENT_CONTEXT_CPU_CYCLE_COUNTER = 3;

const int FLUSH_TIMER_SECOND = 1;

const ULONGLONG KERNEL_KEYWORD_PROCESS = 0x10;

const int LOGGER_NAME_MAX_LENGTH = 20 * MAX_PATH;

std::atomic_bool AppUsingStatistics::isTraceClosed = false;

AppUsingStatistics::AppUsingStatistics() {}

AppUsingStatistics::~AppUsingStatistics() {}

void AppUsingStatistics::Init()
{
    // 初始化白名单
    PROCESSDETECT.Run();
    if (!BuildTraceConfig()) {
        return;
    }
    CloseTrace();
    if (!BuildTraceController(false)) {
        return;
    }
    std::thread([&]() { ConsumeTrace(); }).detach();
    // 应用启动信息消费
    APP_USING_REPORT.Start();
}

void AppUsingStatistics::Quit()
{
    APP_USING_REPORT.Stop();
    PROCESSDETECT.Stop();
    CloseTrace();
    if (pConfig) {
        free(pConfig);
        pConfig = nullptr;
    }
}

bool AppUsingStatistics::BuildTraceConfig()
{
    if (pConfig) {
        return true;
    }
    ULONG uBufferSize = sizeof EVENT_TRACE_PROPERTIES + LOGGER_NAME_MAX_LENGTH * sizeof TCHAR;
    pConfig = (PEVENT_TRACE_PROPERTIES)malloc(uBufferSize);
    if (!pConfig) {
        LOG_ERR("[DFX - BuildTraceConfig] - allocate pConfig failed.");
        return false;
    }
    SecureZeroMemory(pConfig, uBufferSize);
    pConfig->Wnode.BufferSize = uBufferSize;
    pConfig->Wnode.Flags = WNODE_FLAG_TRACED_GUID;
    pConfig->Wnode.ClientContext = CLIENT_CONTEXT_CPU_CYCLE_COUNTER;
    pConfig->LogFileMode = EVENT_TRACE_REAL_TIME_MODE;
    pConfig->LoggerNameOffset = sizeof EVENT_TRACE_PROPERTIES;
    pConfig->FlushTimer = FLUSH_TIMER_SECOND;
    TCHAR *sessionNamePtr = reinterpret_cast<LPSTR>(pConfig) + pConfig->LoggerNameOffset;
    rsize_t size = static_cast<rsize_t>(uBufferSize) - pConfig->LoggerNameOffset;
    if (!sessionNamePtr || strncpy_s(sessionNamePtr, size, SESSION_NAME.c_str(), SESSION_NAME.length())) {
        free(pConfig);
        pConfig = nullptr;
        LOG_ERR("[DFX - BuildTraceConfig] - setting session name failed.");
        return false;
    }
    return true;
}

bool AppUsingStatistics::BuildTraceController(bool isRetry)
{
    if (!pConfig) {
        LOG_ERR("[DFX - BuildTraceController] - pConfig is nullptr.");
        return false;
    }
    TRACEHANDLE traceId = 0;
    HRESULT hres = ::StartTrace(&traceId, SESSION_NAME.c_str(), pConfig);
    if (!traceId) {
        if (hres == ERROR_ACCESS_DENIED) {
            LOG_ERR("[DFX - BuildTraceController] - start trace failed, reason: permission denied.");
        } else if (hres == ERROR_ALREADY_EXISTS) {
            LOG_WARN("[DFX - BuildTraceController] - start trace failed, reason: trace name exists.");
            CloseTrace();
            if (!isRetry) {
                LOG_WARN("[DFX - BuildTraceController] - trying to close trace and start it again.");
                return BuildTraceController(true);
            }
        }
        free(pConfig);
        pConfig = nullptr;
        return false;
    }
    LOG_INFO("[DFX - BuildTraceController] - start trace succeed");
    ULONG isEnabled = 1;
    hres = ::EnableTraceEx(&PsProvGuid, nullptr, traceId, isEnabled, TRACE_LEVEL_NONE, KERNEL_KEYWORD_PROCESS, 0, 0, 0);
    if (hres != ERROR_SUCCESS) {
        LOG_ERR("[DFX - BuildTraceController] - enable trace failed");
        CloseTrace();
        free(pConfig);
        pConfig = nullptr;
        return false;
    }
    return true;
}

int64_t FetchCurrentTimeStamp()
{
    using namespace std::chrono;
    return duration_cast<milliseconds>(system_clock::now().time_since_epoch()).count();
}

VOID WINAPI AppUsingStatistics::ProcessRecordEventsCallback(PEVENT_RECORD pEventRecord)
{
    if (isTraceClosed) {
        return;
    }
    if (!pEventRecord || !pEventRecord->UserData) {
        return;
    }
    AppEventData event;
    switch (pEventRecord->EventHeader.EventDescriptor.Id) {
        case APP_EVENT_TYPE_START:
            event.pid = *(ULONG *)(((PUCHAR)pEventRecord->UserData));
            if (!event.pid) {
                isTraceClosed = true;
                break;
            }
            event.filePath = StringProc::WStringToString(
                std::wstring((LPWSTR)(((PUCHAR)pEventRecord->UserData) + APP_CREATE_EVENT_FILE_PATH_OFFSET)));
            event.name = event.filePath.substr(event.filePath.rfind("\\") + 1);
            event.time = FetchCurrentTimeStamp();
            event.eventType = APP_EVENT_TYPE_START;
            APPCONTROLLER::AppController::OnStartup(event.name);
            PERF_EVENT_DISPATCHER.OnProcessStart(event.name, event.filePath, event.pid);
            PROCESSDETECT.OnRecvProcessMsg(event.pid, ProcessOperation::PROCESS_CREATE);
            PAYLOAD_SENSOR.CompressDecompressApplicationFilter(event.pid, event.name);
            INSTALL_LIFECYCLE_HANDLER.Start(event);
            break;
        case APP_EVENT_TYPE_END:
            event.pid = *(ULONG *)(((PUCHAR)pEventRecord->UserData));
            if (!event.pid) {
                isTraceClosed = true;
                break;
            }
            event.filePath = ((LPSTR)pEventRecord->UserData) + APP_QUIT_EVENT_FILE_PATH_OFFSET;
            event.name = event.filePath.substr(event.filePath.rfind("\\") + 1);
            event.time = FetchCurrentTimeStamp();
            event.eventType = APP_EVENT_TYPE_END;
            PERF_EVENT_DISPATCHER.OnProcessEnd(event.pid);
            PROCESSDETECT.OnRecvProcessMsg(event.pid, ProcessOperation::PROCESS_CLOSE);
            PAYLOAD_SENSOR.CompressDecompressApplicationDeFilter(event.pid);
            break;
        default:
            break;
    }
}

bool AppUsingStatistics::ConsumeTrace()
{
    EVENT_TRACE_LOGFILE logFile;
    SecureZeroMemory(&logFile, sizeof logFile);
    logFile.LoggerName = const_cast<LPSTR>(SESSION_NAME.c_str());
    logFile.EventRecordCallback = ProcessRecordEventsCallback;
    logFile.LogFileMode = EVENT_TRACE_NO_PER_PROCESSOR_BUFFERING | EVENT_TRACE_ADD_HEADER_MODE |
                          EVENT_TRACE_REAL_TIME_MODE;
    hTrace = ::OpenTrace(&logFile);
    if (!hTrace) {
        LOG_ERR("[DFX - ConsumeTrace] - open trace failed!");
        return false;
    }
    isTraceClosed = false;
    if (::ProcessTrace(&hTrace, 1, 0, 0) == ERROR_SUCCESS) {
        LOG_INFO("[DFX - ConsumeTrace] - process trace finished");
        hTrace = 0;
    } else {
        LOG_ERR("[DFX - ConsumeTrace] - process trace failed");
        CloseTrace();
        free(pConfig);
        pConfig = nullptr;
        hTrace = 0;
        return false;
    }
    return true;
}

void AppUsingStatistics::CloseTrace()
{
    if (pConfig) {
        ControlTrace(0, SESSION_NAME.c_str(), pConfig, EVENT_TRACE_CONTROL_STOP);
    }
    if (hTrace) {
        ::CloseTrace(hTrace);
        isTraceClosed = true;
    }
}
}